r""" Evaluate mask prediction """
import cv2 as cv
import numpy as np
import torch
from skimage import morphology


# 评价指标，同论文Beyond the Pixel-Wise Loss for Topology-Aware Delineation中的Pixel-wise metrics章节
# We give precise definitions in
# appendix. In our experiments we use a threshold of 2 pixels
# for roads and cracks, and 1 for the neuronal membranes


# dilate_kernel_size=1，骨架宽度为1，保持不变：用于神经元膜的评估
# dilate_kernel_size=4,骨架宽度从1膨胀到4：对应JTFN论文中we dilate the ground truth masks by 4-pixels in evaluation。
# dilate_kernel_size = 5，骨架宽度从1，膨胀到5，对应论文中In our experiments, the threshold is set to 2，这里threshold指在1个单位像素骨架上下浮动2元素，因此是5的骨架宽度。

class Evaluator():
    def __init__(self, benchmark):
        super(Evaluator, self).__init__()
        if benchmark == 'DRIVE':
            self.dilate_kernel_size = 1
        elif benchmark == 'Crack200':
            self.dilate_kernel_size = 4
        else:
            self.dilate_kernel_size = 5

    # 输入c x h x w （c=1，二值黑白图）
    def evaluateSingleFR(self, pred_mask, gt_mask):  # H x W
        # torch.set_printoptions(profile='full')
        # print('pred',pred_mask)
        # print('gt',gt_mask)
        pred_mask = pred_mask[0].detach().cpu().numpy().astype(int)
        gt_mask = gt_mask[0].detach().cpu().numpy().astype(int)

        tp = np.sum(gt_mask * pred_mask)
        tn = np.sum((1 - gt_mask) * (1 - pred_mask))
        fp = np.sum((1 - gt_mask) * pred_mask)
        fn = np.sum(gt_mask * (1 - pred_mask))
        epsilon = 1e-7
        accuracy = (tp + tn) / (tp + tn + fp + fn + epsilon)
        precision = tp / (tp + fp + epsilon)
        recall = tp / (tp + fn + epsilon)
        f1 = 2 * (precision * recall) / (precision + recall + epsilon)

        return precision * 100, recall * 100, f1 * 100

    # 输入c x h x w （c=1，二值黑白图）
    def evaluateSingleTP(self, pred_mask, gt_mask):
        pred_mask = pred_mask[0].detach().cpu().numpy().astype(int)
        gt_mask = gt_mask[0].detach().cpu().numpy().astype(int)
        # 获取骨架
        pred_skeleton = morphology.skeletonize(pred_mask).astype(float)
        gt_skeleton = morphology.skeletonize(gt_mask).astype(float)

        # 骨架膨胀,单像素骨架
        kernel = cv.getStructuringElement(cv.MORPH_RECT, (self.dilate_kernel_size, self.dilate_kernel_size))
        pred_skeleton = cv.dilate(pred_skeleton, kernel, iterations=1)
        gt_skeleton = cv.dilate(gt_skeleton, kernel, iterations=1)

        epsilon = 1e-7

        # completeness：（标注分割图 ∩ 网络预测分割图）/ 标注分割图
        matched_reference = gt_skeleton * pred_skeleton
        completeness = np.sum(matched_reference) / (np.sum(gt_skeleton) + epsilon)

        # correctness：（网络预测分割图 ∩ 标注分割图）/ 网络预测分割图
        matched_extraction = pred_skeleton * gt_skeleton
        correctness = np.sum(matched_extraction) / (np.sum(pred_skeleton) + epsilon)

        quality = (completeness * correctness) / (completeness - completeness * correctness + correctness + epsilon)

        return correctness * 100, completeness * 100, quality * 100

    # 输入BS x C x H x W ,tensor类型
    # ignore_mask指的是遮罩，例如DRIVE数据集。
    def evaluate(self, pred_mask, gt_mask):
        p, r, f, cor, com, qua = [], [], [], [], [], []

        # each_pred_mask,each_gt_mask: C x H x W
        for each_pred_mask, each_gt_mask in zip(pred_mask, gt_mask):
            precision, recall, f1 = self.evaluateSingleFR(each_pred_mask, each_gt_mask)
            correct, complete, quality = self.evaluateSingleTP(each_pred_mask, each_gt_mask)
            p.append(precision)
            r.append(recall)
            f.append(f1)
            cor.append(correct)
            com.append(complete)
            qua.append(quality)

        evaluates = {
            'Precision': np.average(p),
            'Recall': np.average(r),
            'F1': np.average(f),
            'Correct': np.average(cor),
            'Complete': np.average(com),
            'Quality': np.average(qua),
        }
        # 输出Batchsize均值evaluates
        return evaluates
